// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/connectors/hive/file_handle.h>
#include <pollux/experimental/wave/dwio/ColumnReader.h>
#include <pollux/experimental/wave/exec/tests/utils/FileFormat.h>
#include <pollux/type/subfield.h>

namespace kumo::pollux::wave::test {

class TestFormatData : public wave::FormatData {
 public:
  static constexpr int32_t kNotRegistered = -1;

  TestFormatData(
      OperandId operand,
      int32_t totalRows,
      const test::Column* column)
      : operand_(operand), totalRows_(totalRows), column_(column) {}

  bool hasNulls() const override {
    return column_->nulls != nullptr;
  }

  int32_t totalRows() const override {
    return totalRows_;
  }

  void newBatch(int32_t startRow) override {
    currentRow_ = startRow;
    queued_ = false;
  }

  void griddize(
      ColumnOp& op,
      int32_t blockSize,
      int32_t numBlocks,
      ResultStaging& deviceStaging,
      ResultStaging& resultStaging,
      SplitStaging& staging,
      DecodePrograms& programs,
      ReadStream& stream) override;

  void startOp(
      ColumnOp& op,
      const ColumnOp* previousFilter,
      ResultStaging& deviceStaging,
      ResultStaging& resultStaging,
      SplitStaging& staging,
      DecodePrograms& program,
      ReadStream& stream) override;

 private:
  // Stages movement of nulls to device if any. Returns the id of the buffer or
  // kNotRegisterd.
  int32_t stageNulls(ResultStaging& deviceStaging, SplitStaging& splitStaging);

  void decodeAlphabet(
      ColumnOp& op,
      Column* alphabet,
      ResultStaging& deviceStaging,
      ResultStaging& resultStaging,
      SplitStaging& staging,
      DecodePrograms& programs,
      ReadStream& stream);

  const OperandId operand_;

  int32_t totalRows_{0};

  const test::Column* column_;
  bool staged_{false};
  bool nullsStaged_{false};
  bool queued_{false};
  int32_t numStaged_{0};
  // The device side data area start, set after the staged transfer is done.
  void* deviceBuffer_{nullptr};

  // Device side encoded alphabet for dictionary encoding.
  void* rawAlphabet_{nullptr};

  // Device side decoded alphabet for dictionary encoding.
  void* decodedAlphabet_{nullptr};
  uint32_t* filterBitmap_{nullptr};
};

class TestFormatParams : public wave::FormatParams {
 public:
  TestFormatParams(
      memory::MemoryPool& pool,
      dwio::common::ColumnReaderStatistics& stats,
      const test::Stripe* stripe)
      : FormatParams(pool, stats), stripe_(stripe) {}

  std::unique_ptr<FormatData> toFormatData(
      const std::shared_ptr<const dwio::common::TypeWithId>& type,
      const pollux::common::ScanSpec& scanSpec,
      OperandId operand) override;

  const Stripe* stripe() const {
    return stripe_;
  }

 private:
  const test::Stripe* stripe_;
};

class TestFormatReader {
 public:
  static std::unique_ptr<ColumnReader> build(
      const TypePtr& requestedType,
      const std::shared_ptr<const dwio::common::TypeWithId>& fileType,
      TestFormatParams& params,
      common::ScanSpec& scanSpec,
      std::vector<std::unique_ptr<common::Subfield::PathElement>>& path,
      const DefinesMap& defines,
      bool isRoot = false);
};

} // namespace kumo::pollux::wave::test
